package org.jiucheng.magpiebridge.client.nio;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SocketChannel;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.logging.Level;
import java.util.logging.Logger;

import org.jiucheng.magpiebridge.protocol.Message;
import org.jiucheng.magpiebridge.util.Cfg;

/**
 * 
 * @author jiucheng
 *
 */
public class Client {
    
    private static final Logger LOGGER = Logger.getLogger(Client.class.getName());
    
    private static final ScheduledExecutorService schedule = Executors.newScheduledThreadPool(2);
    private static final ExecutorService executorService = Executors.newCachedThreadPool();
    private static final Map<Integer, SocketChannel> realServers = new ConcurrentHashMap<Integer, SocketChannel>();
    
    // 当前在写
    private static final AtomicBoolean writed = new AtomicBoolean(false);
    private static volatile SocketChannel socketChannel;
    private static volatile boolean runed = false;
    
    public static void main(String[] args) throws IOException, InterruptedException {
        Cfg.loadProperties(Client.class);
        schedule();
        Runtime.getRuntime().addShutdownHook(new Thread() {
            @Override
            public void run() {
                shutdown();
            }
        });
        startup();
    }
    
    private static void schedule() {
        // 重连
        schedule.scheduleAtFixedRate(new Runnable() {
            public void run() {
                if (!runed) {
                    try {
                        startup();
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                }
            }
        }, 60, 60, TimeUnit.SECONDS);
        // 心跳
        schedule.scheduleWithFixedDelay(new Runnable() {
            public void run() {
                SocketChannel sc = Client.socketChannel;
                if (sc != null) {
                    Message message = new Message();
                    message.setMagic(Message.MAGIC);
                    message.setType(Message.Type.HEARTBEAT);
                    ByteBuffer buffer = Message.toByteBuffer(message);
                    while(!writed.compareAndSet(false, true)) {}
                    try {
                        while (buffer.hasRemaining()) {
                            sc.write(buffer);
                        }
                    } catch (IOException e) {
                        e.printStackTrace();
                        closeSocketChannel();
                    } finally {
                        writed.compareAndSet(true, false);
                    }
                }
            }
        }, 30L, 30L, TimeUnit.SECONDS);
    }
    
    public static void startup() {
        runed = true;
        try {
            socketChannel = SocketChannel.open();
            socketChannel.connect(new InetSocketAddress(Cfg.getServerIp(), Cfg.getServerPort()));
        } catch (IOException e) {
            e.printStackTrace();
            closeSocketChannel();
            return;
        }
        
        auth(socketChannel);
        
        ByteBuffer headerByteBuffer = ByteBuffer.allocate(13);
        while (true) {
            headerByteBuffer.clear();
            ByteBuffer buf;
            try {
                int len = socketChannel.read(headerByteBuffer);
                if (len == -1) {
                    closeSocketChannel();
                    return;
                }
                while (headerByteBuffer.position() != headerByteBuffer.capacity()) {
                    socketChannel.read(headerByteBuffer);
                }
                
                headerByteBuffer.flip();
                
                int magic = headerByteBuffer.getInt();
                byte type = headerByteBuffer.get();
                int uri = headerByteBuffer.getInt();
                int size = headerByteBuffer.getInt();
                
                buf = ByteBuffer.allocate(size + 13);
                buf.putInt(magic);
                buf.put(type);
                buf.putInt(uri);
                buf.putInt(size);
                while (buf.position() != buf.capacity()) {
                    socketChannel.read(buf);
                }
            } catch (Exception e) {
                e.printStackTrace();
                closeSocketChannel();
                return;
            }
            
            Message message = Message.fromByteBuffer(buf);
            byte type = message.getType();
            if (type == Message.Type.CONNECT) {
                handleConnectMessage(socketChannel, message);
            } else if (type == Message.Type.DISCONNECT) {
                handleDisconnectMessage(socketChannel, message);
            } else if (type == Message.Type.TRANSFER) {
                handleTransferMessage(socketChannel, message);
            } else if (type == Message.Type.HEARTBEAT) {
                // 心跳不处理
            }
        }
    }
    
    public static void closeSocketChannel() {
        if (socketChannel != null) {
            try {
                socketChannel.close();
            } catch (IOException e) {
                // e.printStackTrace();
            }
            socketChannel = null;
        }
        for (Integer key : realServers.keySet()) {
            SocketChannel socket = realServers.remove(key);
            if (socket != null) {
                try {
                    socket.close();
                } catch (IOException e) {
                    // e.printStackTrace();
                }
            }
        }
        runed = false;
    }
    
    public static void shutdown() {
        schedule.shutdownNow();
        closeSocketChannel();
        if (LOGGER.isLoggable(Level.INFO)) {
            LOGGER.info("shutdown");
        }
    }
    
    private static void handleDisconnectMessage(final SocketChannel socketChannel, final Message proxyMessage) {
        SocketChannel socket = realServers.remove(proxyMessage.getUri());
        if (socket != null) {
            try {
                socket.close();
            } catch (IOException e) {
                // e.printStackTrace();
            }
        }
    }
    
    private static void handleTransferMessage(final SocketChannel socketChannel, final Message message) {
        ByteBuffer byteBuffer = ByteBuffer.allocate(message.getData().length);
        byteBuffer.put(message.getData());
        byteBuffer.flip();
        SocketChannel socket = realServers.get(message.getUri());
        if (socket != null) {
            try {
                while(byteBuffer.hasRemaining()) {
                    socket.write(byteBuffer);
                }
            } catch (Exception e) {
                e.printStackTrace();
                realServers.remove(message.getUri());
                while(!writed.compareAndSet(false, true)) {}
                try {
                    Message rmsg = new Message();
                    rmsg.setMagic(Message.MAGIC);
                    rmsg.setType(Message.Type.DISCONNECT);
                    rmsg.setUri(message.getUri());
                    ByteBuffer buffer = Message.toByteBuffer(rmsg);
                    while(buffer.hasRemaining()) {
                        socketChannel.write(buffer);
                    }
                } catch (Exception e2) {
                    e2.printStackTrace();
                    closeSocketChannel();
                } finally {
                    writed.compareAndSet(true, false);
                }
            }
        }
    }
    
    private static void handleConnectMessage(final SocketChannel socketChannel, final Message message) {
        int uri = message.getUri();
        // lp.io2c.com
        String ipport = new String(message.getData());
        String[] strs = ipport.split(":");
        
        SocketChannel realSocketChannel = null;
        try {
            realSocketChannel = SocketChannel.open();
            realSocketChannel.connect(new InetSocketAddress(strs[0], Integer.valueOf(strs[1])));
        } catch (NumberFormatException e1) {
            e1.printStackTrace();
            disconnect(socketChannel, uri);
            return;
        } catch (IOException e1) {
            e1.printStackTrace();
            disconnect(socketChannel, uri);
            return;
        }
        
        realServers.put(uri, realSocketChannel);
        
        final Integer tmpUri = uri;
        executorService.execute(new Runnable() {
            private Integer uri = tmpUri;
            public void run() {
                    ByteBuffer buf = ByteBuffer.allocate(4 * 1024 * 1024);
                    while(true) {
                        SocketChannel real = realServers.get(uri);
                        if (real == null)  {
                            return;
                        }
                        int len;
                        try {
                            len = real.read(buf);
                        } catch (IOException e) {
                            realServers.remove(uri);
                            disconnect(socketChannel, uri);
                            // e.printStackTrace();
                            return;
                        }
                        if (len == -1) {
                            realServers.remove(uri);
                            disconnect(socketChannel, uri);
                            return;
                        } else if (len > 0) {
                            buf.flip();
                            Message message = new Message();
                            message.setMagic(Message.MAGIC);
                            message.setType(Message.Type.TRANSFER);
                            message.setUri(uri);
                            byte[] bts = new byte[len];
                            buf.get(bts);
                            message.setData(bts);
                            message.setSize(bts.length);
                            ByteBuffer buffer = Message.toByteBuffer(message);
                            while(!writed.compareAndSet(false, true)) {}
                            try {
                                while(buffer.hasRemaining()) {
                                    socketChannel.write(buffer);
                                }
                            } catch (Exception e) {
                                e.printStackTrace();
                                closeSocketChannel();
                                return;
                            } finally {
                                writed.compareAndSet(true, false);
                            }
                            buf.clear();
                        }
                    }
            }
        });
    }
    
    // 代理连接断开消息发送
    private static void disconnect(SocketChannel socketChannel, int uri) {
        Message message = new Message();
        message.setMagic(Message.MAGIC);
        message.setType(Message.Type.DISCONNECT);
        message.setUri(uri);
        ByteBuffer buffer = Message.toByteBuffer(message);
        while (!writed.compareAndSet(false, true)) {}
        try {
            while(buffer.hasRemaining()) {
                socketChannel.write(buffer);
            }
        } catch (Exception e2) {
            e2.printStackTrace();
            closeSocketChannel();
        } finally {
            writed.compareAndSet(true, false);
        }
    }
    
    // 认证信息发送
    private static void auth(final SocketChannel socketChannel) {
        Message message = new Message();
        message.setMagic(Message.MAGIC);
        message.setType(Message.Type.AUTH);
        byte[] data = Cfg.getClientKey().getBytes();
        message.setData(data);
        message.setSize(data.length);
        ByteBuffer buffer = Message.toByteBuffer(message);
        while(!writed.compareAndSet(false, true)) {}
        try {
            while (buffer.hasRemaining()) {
                socketChannel.write(buffer);
            }
        } catch (Exception e) {
            e.printStackTrace();
            closeSocketChannel();
        } finally {
            writed.compareAndSet(true, false);
        }
    }
}
